{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implmentation of LSTM and GRU\n",
    "\n",
    "The maximum accuracy achived on MNIST for both LSTM and GRU is 97%. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning curve for GRU and LSTM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as dsets\n",
    "from torch.autograd import Variable\n",
    "from torch.nn import Parameter\n",
    "from torch import Tensor\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import math\n",
    "\n",
    "\n",
    "cuda = True if torch.cuda.is_available() else False\n",
    "    \n",
    "Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor    \n",
    "\n",
    "torch.manual_seed(125)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(125)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "# from sklearn import datasets\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import TensorDataset\n",
    "\n",
    "\n",
    "class KddData(object):\n",
    "\n",
    "    def __init__(self, batch_size, file_name1, file_name2):\n",
    "        data1 = pd.read_csv(file_name1)\n",
    "        data2 = pd.read_csv(file_name2)\n",
    "        \n",
    "        \n",
    "        self._encoder = {\n",
    "            'label':    LabelEncoder()\n",
    "        }\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        target = np.array(data1['label'])\n",
    "        features = np.array(data1.drop('label', axis=1))\n",
    "        data_X, data_y = self.__encode_data(features, target)\n",
    "        self.train_dataset = TensorDataset(\n",
    "            torch.from_numpy(data_X.astype(np.float32)),\n",
    "            torch.from_numpy(data_y.astype(np.int64))\n",
    "        )\n",
    "\n",
    "        target = np.array(data2['label'])\n",
    "        features = np.array(data2.drop('label', axis=1))\n",
    "        data_X, data_y = self.__encode_data(features, target)\n",
    "        self.test_dataset = TensorDataset(\n",
    "            torch.from_numpy(data_X.astype(np.float32)),\n",
    "            torch.from_numpy(data_y.astype(np.int64))\n",
    "        )\n",
    "\n",
    " \n",
    "        self.train_dataloader = DataLoader(self.train_dataset, self.batch_size, shuffle=True)\n",
    "        self.test_dataloader = DataLoader(self.test_dataset, self.batch_size, shuffle=True)\n",
    "\n",
    "    \"\"\"将数据中字符串部分转换为数字，并将输入的41维特征转换为8*8的矩阵\"\"\"\n",
    "    def __encode_data(self, data_X, data_y):\n",
    "        self._encoder['label'].fit(list(set(data_y)))\n",
    "        data_X = np.pad(data_X, ((0, 0), (0, 100 - len(data_X[0]))), 'constant').reshape(-1, 1, 10, 10)\n",
    "        data_y = self._encoder['label'].transform(data_y)\n",
    "        return data_X, data_y\n",
    "\n",
    "    \"\"\"将数据拆分为训练集和测试集，并转换为TensorDataset对象\"\"\"\n",
    "    def __split_data_to_tensor(self, data_X, data_y):\n",
    "        X_train, X_test, y_train, y_test = train_test_split(data_X, data_y, test_size=0.3)\n",
    "        train_dataset = TensorDataset(\n",
    "            torch.from_numpy(X_train.astype(np.float32)),\n",
    "            torch.from_numpy(y_train.astype(np.int64))\n",
    "        )\n",
    "        test_dataset = TensorDataset(\n",
    "            torch.from_numpy(X_test.astype(np.float32)),\n",
    "            torch.from_numpy(y_test.astype(np.int64))\n",
    "        )\n",
    "        return train_dataset, test_dataset\n",
    "\n",
    "    \"\"\"接受一个数组进行解码\"\"\"\n",
    "    def decode(self, data, label=False):\n",
    "        if not label:\n",
    "            _data = list(data)\n",
    "            _data[1] = self._encoder['protocal'].inverse_transform([_data[1]])[0]\n",
    "            _data[2] = self._encoder['service'].inverse_transform([_data[2]])[0]\n",
    "            _data[2] = self._encoder['flag'].inverse_transform([_data[3]])[0]\n",
    "            return _data\n",
    "        return self._encoder['label'].inverse_transform(data)\n",
    "    \n",
    "    def encode(self, data, label=False):\n",
    "        if not label:\n",
    "            _data = list(data)\n",
    "            _data[1] = self._encoder['protocal'].transform([_data[1]])[0]\n",
    "            _data[2] = self._encoder['service'].transform([_data[2]])[0]\n",
    "            _data[3] = self._encoder['flag'].transform([_data[3]])[0]\n",
    "            return _data\n",
    "        return self._encoder['label'].transform([data])[0]\n",
    "\n",
    "\n",
    "# data_file = 'test-remote-nagle-open-cn2us.csv'\n",
    "# data_file = 'test-cn2cn-nagle-open-wired.csv'\n",
    "# dataset = KddData(batch_size, 'train-expon-nagle-open.csv', data_file)\n",
    "# dataset.train_dataset[0][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "STEP 1: LOADING DATASET\n",
    "'''\n",
    "\n",
    "batch_size = 64\n",
    "\n",
    "trainFile = 'train-wired.csv'\n",
    "testFile = 'test-wired.csv'\n",
    "testFile1 = 'test-kr.csv'\n",
    "testFile2 = 'test-us.csv'\n",
    "testFile3 = 'test-wifi.csv'\n",
    "testFile4 = 'test-4G.csv'\n",
    "testFile5 = 'test-3G.csv'\n",
    "\n",
    "\n",
    "dataset = KddData(batch_size, trainFile, testFile)\n",
    " \n",
    "n_iters = 12000\n",
    "num_epochs = n_iters / (len(dataset.train_dataset) / batch_size)\n",
    "num_epochs = int(num_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = dataset.train_dataloader\n",
    " \n",
    "test_loader = dataset.test_dataloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LSTM cell implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMCell(nn.Module):\n",
    "\n",
    "    \"\"\"\n",
    "    An implementation of Hochreiter & Schmidhuber:\n",
    "    'Long-Short Term Memory' cell.\n",
    "    http://www.bioinf.jku.at/publications/older/2604.pdf\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_size, hidden_size, bias=True):\n",
    "        super(LSTMCell, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.bias = bias\n",
    "        self.x2h = nn.Linear(input_size, 4 * hidden_size, bias=bias)\n",
    "        self.h2h = nn.Linear(hidden_size, 4 * hidden_size, bias=bias)\n",
    "        self.reset_parameters()\n",
    "\n",
    "\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        std = 1.0 / math.sqrt(self.hidden_size)\n",
    "        for w in self.parameters():\n",
    "            w.data.uniform_(-std, std)\n",
    "    \n",
    "    def forward(self, x, hidden):\n",
    "        \n",
    "        hx, cx = hidden\n",
    "        \n",
    "        x = x.view(-1, x.size(1))\n",
    "        \n",
    "        gates = self.x2h(x) + self.h2h(hx)\n",
    "    \n",
    "        gates = gates.squeeze()\n",
    "        \n",
    "        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)\n",
    "        \n",
    "        ingate = F.sigmoid(ingate)\n",
    "        forgetgate = F.sigmoid(forgetgate)\n",
    "        cellgate = F.tanh(cellgate)\n",
    "        outgate = F.sigmoid(outgate)\n",
    "        \n",
    "\n",
    "        cy = torch.mul(cx, forgetgate) +  torch.mul(ingate, cellgate)        \n",
    "\n",
    "        hy = torch.mul(outgate, F.tanh(cy))\n",
    "        \n",
    "        return (hy, cy)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "STEP 3: CREATE MODEL CLASS\n",
    "'''\n",
    " \n",
    "class LSTMModel(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim, bias=True):\n",
    "        super(LSTMModel, self).__init__()\n",
    "        # Hidden dimensions\n",
    "        self.hidden_dim = hidden_dim\n",
    "         \n",
    "        # Number of hidden layers\n",
    "        self.layer_dim = layer_dim\n",
    "               \n",
    "        self.lstm = LSTMCell(input_dim, hidden_dim, layer_dim)  \n",
    "        \n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "     \n",
    "    \n",
    "    \n",
    "    def forward(self, x):\n",
    "        \n",
    "        # Initialize hidden state with zeros\n",
    "        #######################\n",
    "        #  USE GPU FOR MODEL  #\n",
    "        #######################\n",
    "        #print(x.shape,\"x.shape\")100, 28, 28\n",
    "        if torch.cuda.is_available():\n",
    "            h0 = Variable(torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).cuda())\n",
    "        else:\n",
    "            h0 = Variable(torch.zeros(self.layer_dim, x.size(0), self.hidden_dim))\n",
    "\n",
    "        # Initialize cell state\n",
    "        if torch.cuda.is_available():\n",
    "            c0 = Variable(torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).cuda())\n",
    "        else:\n",
    "            c0 = Variable(torch.zeros(self.layer_dim, x.size(0), hidden_dim))\n",
    "\n",
    "                    \n",
    "       \n",
    "        outs = []\n",
    "        \n",
    "        cn = c0[0,:,:]\n",
    "        hn = h0[0,:,:]\n",
    "\n",
    "        for seq in range(x.size(1)):\n",
    "            hn, cn = self.lstm(x[:,seq,:], (hn,cn)) \n",
    "            outs.append(hn)\n",
    "            \n",
    "    \n",
    "\n",
    "        out = outs[-1].squeeze()\n",
    "        \n",
    "        out = self.fc(out) \n",
    "        # out.size() --> 100, 10\n",
    "        return out\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "input_dim = 10\n",
    "hidden_dim = 128\n",
    "layer_dim = 1  # ONLY CHANGE IS HERE FROM ONE LAYER TO TWO LAYER\n",
    "output_dim = 2\n",
    " \n",
    "model = LSTMModel(input_dim, hidden_dim, layer_dim, output_dim)\n",
    "# model = GRUModel(input_dim, hidden_dim, layer_dim, output_dim)\n",
    "\n",
    "#######################\n",
    "#  USE GPU FOR MODEL  #\n",
    "#######################\n",
    " \n",
    "if torch.cuda.is_available():\n",
    "    model.cuda()\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "learning_rate = 0.1\n",
    " \n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LSTM local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Number of steps to unroll\n",
    "seq_dim = 10 \n",
    "\n",
    "loss_list = []\n",
    "iter = 0\n",
    "for epoch in range(num_epochs):\n",
    "    for i, (images, labels) in enumerate(train_loader):\n",
    "        # Load images as Variable\n",
    "        #######################\n",
    "        #  USE GPU FOR MODEL  #\n",
    "        #######################\n",
    "          \n",
    "        if torch.cuda.is_available():\n",
    "            images = Variable(images.view(-1, seq_dim, input_dim).cuda())\n",
    "            labels = Variable(labels.cuda())\n",
    "        else:\n",
    "            images = Variable(images.view(-1, seq_dim, input_dim))\n",
    "            labels = Variable(labels)\n",
    "          \n",
    "        # Clear gradients w.r.t. parameters\n",
    "        optimizer.zero_grad()\n",
    "         \n",
    "        # Forward pass to get output/logits\n",
    "        # outputs.size() --> 100, 10\n",
    "        outputs = model(images)\n",
    "\n",
    "        # Calculate Loss: softmax --> cross entropy loss\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            loss.cuda()\n",
    "\n",
    "        # Getting gradients w.r.t. parameters\n",
    "        loss.backward()\n",
    "\n",
    "        # Updating parameters\n",
    "        optimizer.step()\n",
    "        \n",
    "        loss_list.append(loss.item())\n",
    "        iter += 1\n",
    "         \n",
    "        if iter % 500 == 0:\n",
    "            # Calculate Accuracy         \n",
    "            correct = 0\n",
    "            total = 0\n",
    "            # Iterate through test dataset\n",
    "            for images, labels in test_loader:\n",
    "                #######################\n",
    "                #  USE GPU FOR MODEL  #\n",
    "                #######################\n",
    "                if torch.cuda.is_available():\n",
    "                    images = Variable(images.view(-1, seq_dim, input_dim).cuda())\n",
    "                else:\n",
    "                    images = Variable(images.view(-1 , seq_dim, input_dim))\n",
    "                \n",
    "                # Forward pass only to get logits/output\n",
    "                outputs = model(images)\n",
    "                \n",
    "                # Get predictions from the maximum value\n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                 \n",
    "                # Total number of labels\n",
    "                total += labels.size(0)\n",
    "                 \n",
    "                # Total correct predictions\n",
    "                #######################\n",
    "                #  USE GPU FOR MODEL  #\n",
    "                #######################\n",
    "                if torch.cuda.is_available():\n",
    "                    correct += (predicted.cpu() == labels.cpu()).sum()\n",
    "                else:\n",
    "                    correct += (predicted == labels).sum()\n",
    "             \n",
    "            accuracy = 100 * correct / total\n",
    "             "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5 (default, Sep  3 2020, 21:29:08) [MSC v.1916 64 bit (AMD64)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "e42634819b8c191a5d07eaf23810ff32516dd8d3875f28ec3e488928fbd3c187"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
